package com.buddy.summary.code.ratelimite;

import com.google.common.util.concurrent.RateLimiter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Google的Guava限流测试
 */
public class GuavaRateLimiter {

    private static final Logger logger = LoggerFactory.getLogger(GuavaRateLimiter.class);


    /**
     * 测试方法，每s生产2个令牌，执行结果为500ms打印一次
     */
    public void limiterTest(){
        SimpleDateFormat simpleDateFormat = new SimpleDateFormat("yyyyMMdd HH:mm:ss.SSS");
        RateLimiter rateLimiter = RateLimiter.create(2);
        while(true) {
            rateLimiter.acquire();
            logger.info(simpleDateFormat.format(new Date()));
        }
    }

    private static final ConcurrentHashMap<String,RateLimiter> resourceLimiterMap = new ConcurrentHashMap<>();

    /**
     * 初始化令牌桶
     * @param resource
     * @param qps
     */
    public static void updateResourceQPS(String resource,double qps){
        RateLimiter rateLimiter = resourceLimiterMap.get(resource);
        if (rateLimiter == null){
            rateLimiter = RateLimiter.create(qps);
            //putIfAbsent:如果map中存在当前相同的键值，那么就把value赋值给当前的value，如果没有就进行存储
            RateLimiter putByOtherThread = resourceLimiterMap.putIfAbsent(resource,rateLimiter);
            if (putByOtherThread != null){
                rateLimiter = putByOtherThread;
            }
            rateLimiter.setRate(qps);
        }
    }

    /**
     * 尝试获得令牌
     * @param resource
     */
    public static void tryAcquire(String resource){
        RateLimiter rateLimiter = resourceLimiterMap.get(resource);
        if (rateLimiter == null){
            return;
        }
        if (!rateLimiter.tryAcquire()){
            throw new RuntimeException(resource+"接口访问太频繁");
        }
    }

    /**
     * 移除令牌
     * @param resource
     */
    public static void removeResource(String resource){
        resourceLimiterMap.remove(resource);
    }

}
